package com.qidafang.demo.client;

import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.spec.McpSchema.*;
import io.modelcontextprotocol.spec.McpTransport;
import io.modelcontextprotocol.spec.McpClientTransport;

import java.net.URI;
import java.time.Duration;
import java.util.List;
import java.util.Map;

/**
 * 基于SSE的MCP客户端，连接到指定URL的MCP服务器
 */
public class MySseClient {
    public static void main(String[] args) {
        // 设置控制台输出编码
        System.setProperty("file.encoding", "UTF-8");

        try {
            // 初始化客户端
            McpSyncClient client = initializeClient();
            // 列出可用工具
            listTools(client);
            // 调用工具
            callTool(client);
            // 列出可用资源
            // listResources(client);
            // 读取资源
            // readResource(client);
            // 列出可用提示
            // listPrompts(client);
            // 使用问候提示
            // useGreetingPrompt(client);
            // 关闭客户端
            closeClient(client);
        } catch (Exception e) {
            System.err.println("发生错误: " + e.getMessage());
            e.printStackTrace();
        }
    }

    // 初始化客户端
    private static McpSyncClient initializeClient() {
        try {
            System.out.println("正在初始化MCP客户端...");
            
            // 创建SSE客户端传输层
            // 使用Java内置的HttpClient连接到MCP服务器
            McpClientTransport transport = new HttpClientSseClientTransport("http://192.168.31.100:8045");
            
            // 创建同步客户端
            McpSyncClient client = McpClient.sync(transport)
                .requestTimeout(Duration.ofSeconds(30))  // 增加超时时间
                .capabilities(ClientCapabilities.builder()
                    .roots(true)      // 启用根目录支持
                    .sampling()       // 启用采样支持
                    .build())
                .sampling(request -> {
                    // 简单的采样处理器实现
                    System.out.println("收到采样请求: " + request);
                    return CreateMessageResult.builder()
                        .role(Role.ASSISTANT)
                        .message("这是一个示例响应")
                        .build();
                })
                .build();

            // 初始化连接
            System.out.println("正在连接到MCP服务器...");
            client.initialize();
            System.out.println("MCP客户端初始化完成");
            return client;
        } catch (Exception e) {
            e.printStackTrace();
            throw new RuntimeException("初始化客户端失败", e);
        }
    }

    // 列出可用工具
    private static void listTools(McpSyncClient client) {
        System.out.println("\n获取可用工具列表:");
        ListToolsResult tools = client.listTools();
        tools.tools().forEach(tool -> {
            System.out.println("工具名称: " + tool.name());
            System.out.println("工具描述: " + tool.description());
            System.out.println("工具Schema: " + tool.inputSchema());
            System.out.println("-------------------");
        });
    }

    // 调用计算器工具示例
    private static void callTool(McpSyncClient client) {
        System.out.println("\n调用工具:");
        CallToolResult calcResult = client.callTool(
            new CallToolRequest("searchPerson", 
                Map.of("name", "张飞"))
        );
        System.out.println("调用结果: ");
        calcResult.content().forEach(content -> {
            if (content instanceof TextContent) {
                System.out.println(((TextContent) content).text());
            }
        });
    }

    // 列出可用资源
    private static void listResources(McpSyncClient client) {
        System.out.println("\n获取可用资源列表:");
        ListResourcesResult resources = client.listResources();
        resources.resources().forEach(resource -> {
            System.out.println("资源URI: " + resource.uri());
            System.out.println("资源名称: " + resource.name());
            System.out.println("资源描述: " + resource.description());
            System.out.println("资源类型: " + resource.mimeType());
            System.out.println("-------------------");
        });
    }

    // 读取资源示例
    private static void readResource(McpSyncClient client) {
        ListResourcesResult resources = client.listResources();
        if (!resources.resources().isEmpty()) {
            Resource firstResource = resources.resources().get(0);
            System.out.println("\n读取资源: " + firstResource.uri());
            ReadResourceResult resourceResult = client.readResource(
                new ReadResourceRequest(firstResource.uri())
            );
            System.out.println("资源内容: ");
            resourceResult.contents().forEach(content -> {
                if (content instanceof TextResourceContents) {
                    System.out.println(((TextResourceContents) content).text());
                }
            });
        }
    }

    // 列出可用提示
    private static void listPrompts(McpSyncClient client) {
        System.out.println("\n获取可用提示列表:");
        ListPromptsResult prompts = client.listPrompts();
        prompts.prompts().forEach(prompt -> {
            System.out.println("提示名称: " + prompt.name());
            System.out.println("提示描述: " + prompt.description());
            System.out.println("提示参数: ");
            prompt.arguments().forEach(arg -> {
                System.out.println("  - " + arg.name() + (arg.required() ? " (必需)" : " (可选)") + ": " + arg.description());
            });
            System.out.println("-------------------");
        });
    }

    // 使用提示示例
    private static void useGreetingPrompt(McpSyncClient client) {
        if (client.listPrompts().prompts().stream().anyMatch(prompt -> prompt.name().equals("greeting"))) {
            System.out.println("\n使用问候提示:");
            GetPromptResult promptResult = client.getPrompt(
                new GetPromptRequest("greeting", Map.of("name", "张三"))
            );
            System.out.println("提示结果: " + promptResult.description());
            System.out.println("消息: ");
            promptResult.messages().forEach(message -> {
                System.out.println("角色: " + message.role());
                if (message.content() instanceof TextContent) {
                    System.out.println("内容: " + ((TextContent) message.content()).text());
                }
                System.out.println("-------------------");
            });
        }
    }

    // 关闭客户端
    private static void closeClient(McpSyncClient client) {
        System.out.println("\n正在关闭客户端...");
        client.closeGracefully();
        System.out.println("客户端已关闭");
    }
} 